import torch
import numpy as np

def project_pseudo_lidar_to_rectcam(pts_3d):
    xs, ys, zs = pts_3d[..., 0], pts_3d[..., 1], pts_3d[..., 2]
    return torch.stack([-ys, -zs, xs], dim=-1)

def project_rectcam_to_pseudo_lidar(pts_3d):
    xs, ys, zs = pts_3d[..., 0], pts_3d[..., 1], pts_3d[..., 2]
    return torch.stack([zs, -xs, -ys], dim=-1)

def project_rect_to_image(pts_3d_rect, P, pose_transform=None):
    n = pts_3d_rect.shape[0]
    ones = torch.ones((n, 1), device=pts_3d_rect.device)
    pts_3d_rect = torch.cat([pts_3d_rect, ones], dim=1)
    
    if pose_transform is not None:
        pts_3d_rect = torch.matmul(pts_3d_rect, pose_transform.T)

    pts_2d = torch.mm(pts_3d_rect, torch.transpose(P, 0, 1))  # nx3
    pts_2d[:, 0] /= pts_2d[:, 2]
    pts_2d[:, 1] /= pts_2d[:, 2]
    return pts_2d[:, 0:2]

def unproject_image_to_rect(pts_image, P):
    pts_3d = torch.cat([pts_image[..., :2], torch.ones_like(pts_image[..., 2:3])], -1)
    pts_3d = pts_3d * pts_image[..., 2:3]
    pts_3d = torch.cat([pts_3d, torch.ones_like(pts_3d[..., 2:3])], -1)
    P4x4 = torch.eye(4, dtype=P.dtype, device=P.device)
    P4x4[:3, :] = P
    invP = torch.inverse(P4x4)
    pts_3d = torch.matmul(pts_3d, torch.transpose(invP, 0, 1))
    return pts_3d[..., :3]

def unproject_image_to_pseudo_lidar(pts_image, P):
    pts_3d = torch.cat([pts_image[..., :2], torch.ones_like(pts_image[..., 2:3])], -1)
    pts_3d = pts_3d * pts_image[..., 2:3]
    pts_3d = torch.cat([pts_3d, torch.ones_like(pts_3d[..., 2:3])], -1)
    P4x4 = torch.eye(4, dtype=P.dtype, device=P.device)
    P4x4[:3, :] = P
    invP = torch.inverse(P4x4)
    pts_3d = torch.matmul(pts_3d, torch.transpose(invP, 0, 1))
    pts_3d = pts_3d[..., [2, 0, 1]] * torch.as_tensor([1., -1., -1.], device='cuda')
    return pts_3d


# if __name__ == "__main__":
#     point_cloud_range = np.array([2, -30.4, -3, 59.6, 30.4, 1])
#     voxel_size = [0.2, 0.2, 0.2]
#     grid_size = (
#             point_cloud_range[3:6] - point_cloud_range[0:3]) / np.array(voxel_size)
#     grid_size = np.round(grid_size).astype(np.int64)
#     X_MIN, Y_MIN, Z_MIN = point_cloud_range[:3]
#     X_MAX, Y_MAX, Z_MAX = point_cloud_range[3:]
#     GRID_X_SIZE, GRID_Y_SIZE, GRID_Z_SIZE = grid_size.tolist()
#     VOXEL_X_SIZE, VOXEL_Y_SIZE, VOXEL_Z_SIZE = voxel_size
#     zs = torch.linspace(Z_MIN + VOXEL_Z_SIZE / 2., Z_MAX - VOXEL_Z_SIZE / 2.,
#                         GRID_Z_SIZE, dtype=torch.float32)
#     ys = torch.linspace(Y_MIN + VOXEL_Y_SIZE / 2., Y_MAX - VOXEL_Y_SIZE / 2.,
#                         GRID_Y_SIZE, dtype=torch.float32)
#     xs = torch.linspace(X_MIN + VOXEL_X_SIZE / 2., X_MAX - VOXEL_X_SIZE / 2.,
#                         GRID_X_SIZE, dtype=torch.float32)
#     zs, ys, xs = torch.meshgrid(zs, ys, xs)
#     coordinates_3d = torch.stack([xs, ys, zs], dim=-1).float()
#     coordinates_3d = coordinates_3d.reshape(-1, 3)
#     rec_c3d = project_pseudo_lidar_to_rectcam(coordinates_3d)
#     contents = [l.strip("\n") for l in open("/cv/yc/DSGN2/data/ww/training/calib/1692074207250000.txt", "r").readlines()]
#     contents = np.array([l.split(' ') for l in contents])
#     import pdb
#     pdb.set_trace()
#     p2 = np.array(contents[0][1:]).astype(float).reshape(3,4)
#     p2 = torch.tensor(p2, dtype=torch.float32)
#     rec2img = project_rect_to_image(rec_c3d, p2)
#     print(rec2img.shape)